#ifndef CKKS_SRC_CONTEXT_H_
#define CKKS_SRC_CONTEXT_H_

#include <complex>
#include <chrono>
#include <map>
#include <algorithm>
#include <random>

#include <arm_neon.h>
//#include <cstdint>


#include <NTL/ZZ.h>
using namespace NTL;

#include "common.h"
#include "numb.h"

using namespace std;

using uint128_t = unsigned __int128;

// 模数参数 logN15QP891为算子测试，logN12QP1725与logN15QP1725为针对运行自举的参数
enum class params { logN15QP891 = 0, logN12QP1725 = 1, logN15QP881 = 2, logN15QP1725 = 3 };


class Context {
public:
    // Encryption parameters
    long logN;  // /< Logarithm of Ring Dimension
    int lambda;
    long logNh; // /< Logarithm of Ring Dimension - 1
    long L;     // /< Maximum Level that we want to support
    long K;     // /< The number of special modulus (usually L + 1)
    long N;
    long M;
    long Nh;

    long slots;

    static uint64_t BASECONVERT;
    
    complex<double> *uVals;
    uint64_t *tmpCode;
    uint64_t *raRaiseAndEqual;
    uint64_t *raReScale;
    uint64_t *raBackAndEqual;
    uint64_t *raReScaleAndEqual;
    uint64_t *tmpLeftRot;

    uint64_t *tmp3;
    uint64_t *tmpDecodeSingle;

    long logp;
    __uint128_t p;

    long h;
    double sigma;

    long T;
    long KN;
    long TN;
    int gamma; // 模切片
    long blockLen;
    int dNum;

    uint64_t *polyQPool;
    uint64_t **polyTBootStrap;
    uint64_t *halfTPolyT;
    uint64_t *halfTPolyPQ;
    uint64_t *halfTPolyPi;
    uint64_t *halfTPolyQi;
    uint64_t *halfTPolyti;

    uint64_t *ex;
    uint64_t *sxSx;
    uint64_t *sxConj;
    vector<vector<uint64_t>> resultax;
    vector<vector<uint64_t>> resultbx;
    vector<vector<vector<uint64_t>>> rax;
    vector<vector<vector<uint64_t>>> rbx;

    vector<uint64_t *> polyTPools;
    vector<uint64_t> PQ;
    vector<uint64_t> QP;
    vector<vector<uint64_t>> PQSlices;
    vector<vector<uint64_t>> PQSlicesHatModPQ;
    vector<vector<uint64_t>> PQSlicesHatInvModPQ;
    vector<vector<uint64_t>> PQSlicesHatInvModPQShoup;
    vector<vector<vector<uint64_t>>> PQSlicesHatModt;
    vector<vector<uint64_t>> PQSlicesModt;

    vector<vector<vector<uint64_t>>> qSlicesHatModq;
    vector<vector<vector<uint64_t>>> qSlicesHatInvModq;
    vector<vector<vector<uint64_t>>> qSlicesHatInvModqShoup;
    vector<vector<vector<vector<uint64_t>>>> qSlicesHatModt;
    vector<vector<vector<uint64_t>>> QSlicesModt;

    vector<uint64_t> tHatModt;
    vector<uint64_t> tHatInvModt;
    vector<uint64_t> tHatInvModtShoup;
    vector<vector<uint64_t>> tHatModPQ;
    vector<uint64_t> TModPQ;


    vector<uint64_t> qVec;
    vector<uint64_t> pVec;
    vector<uint64_t> tVec;
    float *qVecInv;
    float *tVecInv;
    uint64_t *twoqVec;
    
    uint64_t *twopVec;
    
    uint64_t *twotVec;

    __uint128_t *qrVec;
    __uint128_t *prVec;
    __uint128_t *trVec;

    vector<uint64_t> pqrVecHigh;
    vector<uint64_t> pqrVecLow;

    vector<uint64_t> twoqpVec;
    vector<uint64_t> qprVecHigh;
    vector<uint64_t> qprVecLow;

    uint64_t *qdVec;
    uint64_t *pdVec;
    uint64_t *tdVec;

    uint64_t *qrVecHigh; // Barrett recution
    uint64_t *qrVecLow;  // Barrett recution
    uint64_t *prVecHigh; // Barrett recution
    uint64_t *prVecLow;  // Barrett recution
    uint64_t *trVecHigh; // Barrett recution
    uint64_t *trVecLow;  // Barrett recution

    uint64_t *qInvVec;
    uint64_t *pInvVec;
    uint64_t *tInvVec;

    uint64_t *qRoots;
    uint64_t *pRoots;
    uint64_t *tRoots;


    uint64_t *qRootsInv;
    uint64_t *pRootsInv;
    uint64_t *tRootsInv;

    uint64_t **qRootPows;
    uint64_t **pRootPows;
    uint64_t **tRootPows;

    uint64_t **qRootPowsShoup;    // shoup参数
    uint64_t **qRootPowsInvShoup; // shoup参数
    uint64_t **pRootPowsShoup;    // shoup参数
    uint64_t **pRootPowsInvShoup; // shoup参数
    uint64_t **tRootPowsShoup;    // shoup参数
    uint64_t **tRootPowsInvShoup; // shoup参数

    uint64_t **qRootPowsInv;
    uint64_t **pRootPowsInv;
    uint64_t **tRootPowsInv;

    uint64_t *NInvModq;
    uint64_t *NInvModp;
    uint64_t *NInvModt;

    uint64_t *NInvModqShoup; // shoup参数
    uint64_t *NInvModpShoup; // shoup参数
    uint64_t *NInvModtShoup; // shoup参数


    uint64_t **qHatModq; // [l][i] (phat_i)_l mod p_i
    uint64_t *pHatModp;  // [k] qhat_k mod q_k

    uint64_t **qHatInvModq; // [l][i] (qhat_i)_l^-1 mod q_i
    uint64_t *pHatInvModp;  // [k] phat_k^-1 mod p_k
    uint64_t *pHatInvModpShoup;

    uint64_t ***qHatModp; // [l] [i] [k]  (phat_i)_l mod q_k

    uint64_t **pHatModq; // [k][i] qhat_k mod p_i

    uint64_t *PModq;    // [i] qprod mod p_i
    uint64_t *PInvModq; // [i] qprod mod p_i
    uint64_t *PInvModqShoup;

    uint64_t **QModp;    // [i] qprod mod p_i
    uint64_t **QInvModp; // [i] qprod mod p_i

    uint64_t **qInvModq; // [i][j] p_i^-1 mod p_j

    long *rotGroup; // /< precomputed rotation group indexes

    complex<double> *ksiPows; // /< precomputed ksi powers

    map<string, double *> taylorCoeffsMap; // /< precomputed taylor coefficients

    vector<double> mVals; // 改进的采用所用到中间值

    // 构造函数
    Context(long logN, long logp, long L, long K, long T, int gamma, int lambda = 0, int dNum = 0, long h = 64, double sigma = 3.2);

    static Context InitContext(long logN, int lambda = 0, int d = 0);

    static Context InitContextBootStrapping(long logN);

    // 拷贝构造函数
    Context(const Context &other) = delete;
    Context &operator = (const Context &other) = delete;
    
    // 移动构造函数
    Context(Context &&other) = delete;
    Context &operator = (Context &&other) = delete;

    void testparams();

    void GetPrimes(uint64_t lambda, uint64_t logN, uint64_t bit);

    void GetParameters(params param);

    void GetParametersLambda(long logN, int lambda, int dNum);

    void ArrayBitReverse(complex<double> *vals, const long size);

    void ArrayBitReverse(uint64_t *vals, const long size);

    uint64_t FastPow(uint64_t a, uint64_t k, uint64_t mod);

    void FFTSpecial(complex<double> *vals, const long size);
    void FFTSpecialInvLazy(complex<double> *vals, const long size);
    void FFTSpecialInv(complex<double> *vals, const long size);

    void EncodePQ(uint64_t *ax, complex<double> *vals, double scale, long slotsEncode, long l, long k); // CTS中要用到
    void Encode(uint64_t *ax, complex<double> *vals, long slotsEncode, long l);
    void Encode(uint64_t *ax, double *vals, long slots, long l);

    void EncodeSingle(uint64_t *axEncode, complex<double> &val, long l);
    void EncodeSingle(uint64_t *ax, double val, long l);

    void Decode(uint64_t *ax, complex<double> *vals, long slotsDecode, long l, __uint128_t scale);

    void PolyToBigintLvl(int level, uint64_t *p1, int gap, std::vector<ZZ> &coeffsBigint);

    void Decode(uint64_t *ax, double *vals, long slots, long l);

    void DecodeSingle(uint64_t *ax, complex<double> &val);

    void DecodeSingle(uint64_t *ax, double val);

    void QiNTT(uint64_t *res, uint64_t *a, long index);

    void PiNTT(uint64_t *res, uint64_t *a, long index);

    void NTT(uint64_t *res, uint64_t *a, long l, long k = 0);

    void QiNTTAndEqual(uint64_t *a, long index);
    void PiNTTAndEqual(uint64_t *a, long index);
    void TiNTTAndEqual(uint64_t *a, long index);

    void NTTAndEqual(uint64_t *a, long l, long k = 0);
    void TNTTAndEqual(uint64_t *a, long t);

    void QiINTT(uint64_t *res, uint64_t *a, long index);
    void PiINTT(uint64_t *res, uint64_t *a, long index);


    void INTT(uint64_t *res, uint64_t *a, long l, long k = 0);

    void QiINTTAndEqual(uint64_t *a, long index);
    void PiINTTAndEqual(uint64_t *a, long index);
    void TiINTTAndEqual(uint64_t *a, long index);

    void INTTAndEqual(uint64_t *a, long l, long k = 0);
    void TINTTAndEqual(uint64_t *a, long t);

    void QiNegate(uint64_t *res, uint64_t *a, long index);
    void PiNegate(uint64_t *res, uint64_t *a, long index);

    void Negate(uint64_t *res, uint64_t *a, long l, long k = 0);

    void QiNegateAndEqual(uint64_t *a, long index);
    void PiNegateAndEqual(uint64_t *a, long index);

    void NegateAndEqual(uint64_t *a, long l, long k = 0);

    void QiAddConst(uint64_t *res, uint64_t *a, uint64_t c, long index);
    void PiAddConst(uint64_t *res, uint64_t *a, uint64_t c, long index);

    void AddConst(uint64_t *res, uint64_t *a, uint64_t c, long l, long k = 0);

    void QiAddConstAndEqual(uint64_t *a, uint64_t c, long index);
    void PiAddConstAndEqual(uint64_t *a, uint64_t c, long index);

    void AddConstAndEqual(uint64_t *a, uint64_t c, long l, long k = 0);

    void QiSubConst(uint64_t *res, uint64_t *a, uint64_t c, long index);
    void PiSubConst(uint64_t *res, uint64_t *a, uint64_t c, long index);

    void SubConst(uint64_t *res, uint64_t *a, uint64_t c, long l, long k = 0);

    void QiSubConstAndEqual(uint64_t *a, uint64_t c, long index);
    void PiSubConstAndEqual(uint64_t *a, uint64_t c, long index);

    void SubConstAndEqual(uint64_t *a, uint64_t c, long l, long k = 0);

    void QiAdd(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void PiAdd(uint64_t *res, uint64_t *a, uint64_t *b, long index);

    void Add(uint64_t *res, uint64_t *a, uint64_t *b, long l, long k = 0);

    void QiAddAndEqual(uint64_t *a, uint64_t *b, long index);
    void PiAddAndEqual(uint64_t *a, uint64_t *b, long index);
    void TiAddAndEqual(uint64_t *a, uint64_t *b, long index);
    void TAddAndEqual(uint64_t *a, uint64_t *b, long t);

    void AddAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);

    void QiSub(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void PiSub(uint64_t *res, uint64_t *a, uint64_t *b, long index);


    void Sub(uint64_t *res, uint64_t *a, uint64_t *b, long l, long k = 0);


    void QiSubAndEqual(uint64_t *a, uint64_t *b, long index);
    void PiSubAndEqual(uint64_t *a, uint64_t *b, long index);


    void SubAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);
    void TSubAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);

    void QiSub2AndEqual(uint64_t *a, uint64_t *b, long index);
    void PiSub2AndEqual(uint64_t *a, uint64_t *b, long index);

    void Sub2AndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);


    void QiMulConstAndEqual(uint64_t *res, uint64_t cnst, long index);
    void PiMulConstAndEqual(uint64_t *res, uint64_t cnst, long index);

    void MulConstAndEqual(uint64_t *res, uint64_t cnst, long l, long k = 0);

    void QiMul(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void PiMul(uint64_t *res, uint64_t *a, uint64_t *b, long index);

    void Mul(uint64_t *res, uint64_t *a, uint64_t *b, long l, long k = 0);

    void QiMulAndEqual(uint64_t *a, uint64_t *b, long index);
    void PiMulAndEqual(uint64_t *a, uint64_t *b, long index);

    void MulAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);

    void QiSquare(uint64_t *res, uint64_t *a, long index);
    void PiSquare(uint64_t *res, uint64_t *a, long index);

    void Square(uint64_t *res, uint64_t *a, long l, long k = 0);


    void SquareAndEqual(uint64_t *a, long l, long k = 0);

    void EvalAndEqual(uint64_t *a, long l);

    void RaiseAndEqual(uint64_t *&a, long l);
    void Raise(uint64_t *res, uint64_t *a, long l);

    void BackAndEqual(uint64_t *&a, long l);
    void Back(uint64_t *res, uint64_t *a, long l);

    void ReScaleAndEqual(uint64_t *a, long l);

    void ReScale(uint64_t *res, uint64_t *a, long l);

    void ModDownAndEqual(uint64_t *&a, long l, long dl);
    uint64_t *ModDown(uint64_t *a, long l, long dl);

    void LeftRot(uint64_t *res, uint64_t *a, long l, long rotSlots);
    void LeftRotAndEqual(uint64_t *a, long l, long rotSlots);

    void Conjugate(uint64_t *res, uint64_t *a, long l);
    void ConjugateAndEqual(uint64_t *a, long l);

    void SampleGauss(uint64_t *res, long l, long k = 0);
    void SampleZO(uint64_t *res, long l, long k = 0);
    void SampleUniform(uint64_t *res, long l, long k = 0);
    void SampleHWT(uint64_t *res, long l, long k = 0);

    uint64_t BitReverse64(uint64_t index, uint64_t bitLen);
    void PermuteNTTIndex(uint64_t galEl, uint64_t *ntt_index);
    void PermuteNTTWithIndex(uint64_t *polOut, uint64_t *polIn, uint64_t *index, int level);
    void IPermuteNTTWithIndex(uint64_t *polOut, uint64_t *polIn, uint64_t *index);

    vector<vector<uint64_t>> BaseConvertPQ2T(uint64_t *a);

    void BaseConvertQl2T(uint64_t *a, long l, vector<uint64_t *> &res);

    vector<uint64_t *> BaseConvertQl2T(uint64_t *a, long l);
    
    uint64_t *SlicesBaseConvertQl2T(uint64_t *a, long length, long slices, long offset);

    void BaseConvertT2PQl(uint64_t *res, uint64_t **a, long l);

    void BaseConvertPQl2Ql(uint64_t *res, uint64_t *a, long l);
};

#endif
